# central settings
#
simulation     = False
year           = "2018"
magnetPolarity = "Down"
###


########################################################################
from os import environ
import GaudiKernel.SystemOfUnits as Units
from Gaudi.Configuration import *
from Configurables import (GaudiSequencer,
                           DecayTreeTuple,
                           CombineParticles,
                           LoKi__Hybrid__TupleTool,
                           DaVinci,
                           FilterDesktop)

from Configurables import LoKi__Hybrid__PlotTool as PlotTool
from Configurables import LoKi__Hybrid__FilterCriterion as LoKiFilterCriterion
from DecayTreeTuple.Configuration import *

from GaudiConfUtils.ConfigurableGenerators import  CombineParticles
from PhysSelPython.Wrappers import (SelectionSequence, Selection, DataOnDemand, AutomaticData, MergedSelection)
from Configurables import LoKi__VoidFilter as Filter
from PhysConf.Filters import LoKi_Filters
LoKiTool = LoKi__Hybrid__TupleTool( 'LoKiTool')
from Configurables import TriggerTisTos
#
######################################################################################################

if simulation is True:
    location = 'AllStreams/Phys/FullDSTDiMuonJpsi2MuMuDetachedLine/Particles'#MC的TES位置
if simulation is False:
    location = 'Dimuon/Phys/FullDSTDiMuonJpsi2MuMuDetachedLine/Particles'#data的TES位置
######################################################################################################
# if you are setting something wrong, it should die here...
if year not in [ "2017", "2018" ]:
    raise RuntimeError(year + " is not a valid year")
if magnetPolarity not in ["Up", "Down"]:
    raise RuntimeError(magnetPolarity + " is not a valid magnet polarity")
######################################################################################################
# CUTS CUTS CUTS CUTS CUTS CUTS CUTS
####################################
BdCombCuts = "(AM > 4800 * MeV) & "\
            "(AM < 7000 * MeV)"
BdCuts = "(M                                   > 4800 * MeV) & "\
    "(M                                   < 7000 * MeV) &"\
    "(BPVIPCHI2()                         < 16) & "\
    "(BPVDIRA                             > 0.9995) & "\
    "(VFASPF(VCHI2/VDOF)                  < 10)"

Pi0Cuts = "(M > 0 * MeV) &"\
    "(M < 300 * MeV) &"\
    "(PT > 1000 * MeV)"

######################################################################################################
###

def __Pi0__(Pi0Cuts):
    """
    Filter pi0 from StdLooseResolvedPi0
    """
    _pi0resolved = AutomaticData(Location = 'Phys/StdLooseResolvedPi0/Particles')
    _filter_pi0resolved = FilterDesktop("ResolvedPi0Filter", Code = Pi0Cuts )
    _sel = Selection("Selection_pi0resolved",
                                RequiredSelections = [ _pi0resolved ] ,
                                Algorithm = _filter_pi0resolved) 
    return _sel

def __jpsi__(location):
    """
    Get the J/psi from the stripping line.
    Note the different location for data and MC
    """
    return AutomaticData(Location = location)

###
def __jpsipi0__(Jpsi,Pi0, BdCombCuts, BdCuts):
    _jpsipi0 = CombineParticles()
    _jpsipi0.DecayDescriptor = "[B0 -> J/psi(1S) pi0]cc"
    _jpsipi0.CombinationCut = BdCombCuts
    _jpsipi0.MotherCut = BdCuts
    _jpsipi0Conf = _jpsipi0.configurable("Combine_BJpsipi0")
    _selJpsipi0 = Selection("BJpsipi0",
                              Algorithm = _jpsipi0Conf,
                              RequiredSelections = [ Jpsi, Pi0])
    return _selJpsipi0

#######
Pi0             = __Pi0__(Pi0Cuts)
Jpsi            = __jpsi__(location)
Jpsipi0   = __jpsipi0__(Jpsi,Pi0,BdCombCuts, BdCuts)
BJpsipi0Seq = SelectionSequence( 'BJpsipi0Seq', TopSelection = Jpsipi0 )

######################################################################################################
BJpsiPi0 = DecayTreeTuple("BJpsiPi0Tuple")

#input location
BJpsiPi0.Inputs = [ Jpsipi0.outputLocation() ]

# decay descriptors
BJpsiPi0.Decay = "[B0 -> ^(J/psi(1S) -> ^mu+ ^mu-) ^(pi0 -> ^gamma ^gamma)]CC"

# define the tools and configure them

BJpsiPi0.ToolList = [
    "TupleToolKinematic"  #运动学变量，如动量、质量等
    ,"TupleToolEventInfo"
    ]


# define the list of triggers that could have fired...
triggerList = [ 
                "L0DiMuonDecision",
                "L0MuonDecision",
                "Hlt1TrackMuonDecision",
                "Hlt1DiMuonLowMassDecision",
                "Hlt1DiMuonHighMassDecision",
                "Hlt1SingleMuonHighPTDecision",
                "Hlt2DiMuonDetachedJPsiDecision",
                "Hlt2DiMuonDetachedHeavyDecision" ]

# define the tools and configure them
BJpsiPi0.addTupleTool("LoKi::Hybrid::TupleTool/LoKiTool")
BJpsiPi0.LoKiTool.Variables = {
    "eta"                     : "ETA",
    "phi"                     : "PHI",
    "LV01"                    : "LV01",
    "BPVDLS"                  : "BPVDLS" }


# some truthmatching stuff
if simulation is True:
    MCTruth=BJpsiPi0.addTupleTool("TupleToolMCTruth/MCTruth")
    MCTruth.addTupleTool("MCTupleToolHierarchy")

####

####

################################################################
################################################################
# define the branches for resolved pi0
# for resolved pi0
BJpsiPi0.addBranches({  # remove all "^" except where needed.
        "B"          : "^[B0 -> (J/psi(1S) ->  mu+  mu-) ( pi0 ->  gamma  gamma)]CC",
        "mu1"     : " [B0 -> (J/psi(1S) ->  mu+ ^mu-) ( pi0 ->  gamma  gamma)]CC",
        "mu2"    : " [B0 -> (J/psi(1S) -> ^mu+  mu-) ( pi0 ->  gamma  gamma)]CC",
        "Jpsi"       : " [B0 -> ^(J/psi(1S) -> mu+  mu-) ( pi0 ->  gamma  gamma)]CC",
        "pi0": " [B0 -> (J/psi(1S) ->  mu+  mu-) ^(pi0 ->  gamma  gamma)]CC",
        "gamma0"     : " [B0 -> (J/psi(1S) ->  mu+  mu-) ( pi0 -> ^gamma  gamma)]CC",
        "gamma1"     : " [B0 -> (J/psi(1S) ->  mu+  mu-) ( pi0 ->  gamma ^gamma)]CC"
        })

### B tools
BJpsiPi0.B.addTupleTool("TupleToolDecayTreeFitter/BDTF")
BJpsiPi0.B.BDTF.constrainToOriginVertex = True
BJpsiPi0.B.BDTF.daughtersToConstrain = [ 'J/psi(1S)', 'pi0' ]
BJpsiPi0.B.addTupleTool("TupleToolGeometry")#粒子的原初顶点/衰变顶点等的卡方值，飞行距离等。
BJpsiPi0.B.addTupleTool("TupleToolPropertime")#这个工具用于计算粒子的寿命信息，通常是通过测量从产生到衰变的时间间隔来计算的。
BJpsiPi0.B.addTupleTool("TupleToolVtxIsoln")#这个工具用于评估粒子顶点的隔离度，即粒子顶点周围有多少其他粒子活动，这有助于判断粒子是否可能是从一个较干净的衰变过程中产生。
BJpsiPi0.B.addTupleTool("TupleToolConeIsolation/TupleToolConeIsolation")#这个工具用于计算锥形隔离度
BJpsiPi0.B.TupleToolConeIsolation.MinConeSize = 1.8
BJpsiPi0.B.TupleToolConeIsolation.MaxConeSize = 1.8
BJpsiPi0.B.TupleToolConeIsolation.FillComponents = True #填充所有组分信息
BJpsiPi0.B.TupleToolConeIsolation.FillNeutral = True #不填充中性锥角
BJpsiPi0.B.TupleToolConeIsolation.FillMaxPt = False #不填充锥角最大动量信息
BJpsiPi0.B.TupleToolConeIsolation.FillDeltas = True
BJpsiPi0.B.TupleToolConeIsolation.FillAsymmetry = True


### Jpsi tools
BJpsiPi0.Jpsi.addTupleTool("TupleToolGeometry")
BJpsiPi0.Jpsi.addTupleTool( "TupleToolTISTOS/TupleToolTISTOS")
BJpsiPi0.Jpsi.TupleToolTISTOS.TriggerList = triggerList
BJpsiPi0.Jpsi.TupleToolTISTOS.VerboseL0 = True
BJpsiPi0.Jpsi.TupleToolTISTOS.VerboseHlt1 = True
BJpsiPi0.Jpsi.TupleToolTISTOS.VerboseHlt2 = True
BJpsiPi0.Jpsi.addTupleTool("TupleToolAngles")

# mu1tools
BJpsiPi0.mu1.addTupleTool("TupleToolGeometry")
BJpsiPi0.mu1.addTupleTool("TupleToolPid")#ProbNNmu误判别等信息

# mu2 tools
BJpsiPi0.mu2.addTupleTool("TupleToolGeometry")
BJpsiPi0.mu2.addTupleTool("TupleToolPid")


# pi0 tools
BJpsiPi0.pi0.addTupleTool("TupleToolPi0Info")#用于提取或计算与pi0粒子相关的信息的工具。
BJpsiPi0.pi0.addTupleTool("TupleToolAngles")

# gamma tools
BJpsiPi0.gamma0.addTupleTool("TupleToolPhotonInfo")
BJpsiPi0.gamma0.addTupleTool("TupleToolProtoPData")
BJpsiPi0.gamma0.TupleToolProtoPData.DataList += ["IsNotH", "IsNotE", "IsPhoton", "CaloNeutralSpd","CaloNeutralID"]
#TupleToolCaloHypo:通常用于处理和分析与量能器（Calorimeter）相关的假设（Hypothesis）数据.


# gamma tools
BJpsiPi0.gamma1.addTupleTool("TupleToolPhotonInfo")
BJpsiPi0.gamma1.addTupleTool("TupleToolProtoPData") #Prototype Particle Data：意为原型粒子数据、指的是这些实验数据中的原始或原型粒子信息。
BJpsiPi0.gamma1.TupleToolProtoPData.DataList += ["IsNotH", "IsNotE", "IsPhoton", "CaloNeutralSpd","CaloNeutralID"]
#CaloNeutralSpd、CaloNeutralID
#Calorimeter的缩写，表示量能器、Neutral"：表示中性粒子、Spd"：是"Scintillating Pad Detector"的缩写，表示闪烁垫探测器、ID"：是"Inner Detector"的缩写，表示内部探测器
# 在量能器中测量中性粒子在探测器中的能量信息


################################################################
################################################################
Seq1 = GaudiSequencer("Seq1")
Seq1.Members.append( BJpsipi0Seq.sequence() )
Seq1.Members.append( BJpsiPi0 )
Seq1.ModeOR = True
Seq1.ShortCircuit = False
# ######################################################################################################
DaVinci().MainOptions = ""
DaVinci().UserAlgorithms = [Seq1]
# ######################################################################################################
stripFilter = LoKi_Filters (
    STRIP_Code = " HLT_PASS_RE ('StrippingFullDSTDiMuonJpsi2MuMuDetachedLineDecision') "
)
stripFilterSeq = stripFilter.sequence( "StripFilter" )
########################################################################
DaVinci().EvtMax    = -1
DaVinci().DataType  = year
DaVinci().PrintFreq = 1000
DaVinci().EventPreFilters = [ stripFilterSeq ]
DaVinci().InputType = "DST"

if simulation is True:
    DaVinci().Simulation = True
    DaVinci().Lumi       = False
    DaVinci().TupleFile  = "Tuple.root"

if simulation is False:
    DaVinci().TupleFile = "Tuple.root"

# db tags                                                                                                                                             
if simulation is True:
    if year == "2017":
        if magnetPolarity == "Down":
            DaVinci().CondDBtag = "sim-20130522-vc-md100"
            DaVinci().DDDBtag   = "dddb-20130929"
        if magnetPolarity == "Up":
            DaVinci().CondDBtag = "sim-20130522-vc-mu100"
            DaVinci().DDDBtag   = "dddb-20130929"

    if year == "2018":
        if magnetPolarity == "Down":
            DaVinci().CondDBtag = "sim-20190430-vc-md100"
            DaVinci().DDDBtag   = "dddb-20170721-3"
        if magnetPolarity == "Up":
            DaVinci().CondDBtag = "sim-20190430-vc-mu100"
            DaVinci().DDDBtag   = "dddb-20170721-3"

